# -*- coding: utf-8 -*-
"""
Created on Tue Sep 17 14:36:26 2013

@author: jkwong
"""

#PBAR_ZspecBasic_FirstLook.py

import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os
import matplotlib.pyplot as plt
from scipy import interpolate
from matplotlib import cm
from scipy import ndimage

# list of zpsec detector names - Z1 to Z136
ZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137)])
# list of bad detector numbers starting at 1
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
# list of good detector numbers starting at 1
goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])
# list of bad detector names
badZspecNameList = ZspecNameList[(badZspecList-1)]
# list of good detector names
goodZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecIndices = badZspecList - 1
goodZspecIndices = goodZspecList - 1

badZspecMask = np.zeros(len(ZspecNameList))
badZspecMask[badZspecIndices] = True

goodZspecMask = np.zeros(len(ZspecNameList))
goodZspecMask[goodZspecIndices] = True


#################################
## DEFINE FILE NAMES

# Zspec Files
fullFilenameList = []
fullFilenameList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\4761-FDFC-All2148.npy')
fullFilenameList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\4763-FDFC-All2151.npy')

filenameList = []
for f in fullFilenameList:
    a, b = os.path.split(f)
    filenameList.append(b)

# Corresponding cargo images
filenameCargoList = []
filenameCargoList.append('PBAR-20130902007.cargoimage')
filenameCargoList.append('PBAR-20130902009.cargoimage')

fullFilenameCargoList = []
fullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\PBAR-20130902007.cargoimage')
fullFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\PBAR-20130902009.cargoimage')

# Marker files
markerFilenameCargoList = []
markerFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4761\PBAR-20130902007.cargomarker')
markerFilenameCargoList.append(r'C:\Users\jkwong\Documents\Work\PBAR\data3\4763\PBAR-20130902009.cargomarker')

# some number related to the datasets
acquisitionTime = np.array([1/60., 1/60.])
airSection = np.array([0, 175])

numberTimeSlices = 2000
preCargoTimeSlices = 100

######################
## READ IN DATA

# Read in marker files
markerList = []
for (index, f) in enumerate(markerFilenameCargoList):
    markerList.append(PBAR_Cargo.ReadCargoMarker(f))

# read in cargo images
datCargo = []
startBinCargoList = []
datCargoStandardWidth = []
for (index, f) in enumerate(fullFilenameCargoList):
    (A,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(f)
    datCargo.append(A)
    # flip, cut to 2000, flip again    
#    datCargoStandardWidth.append(A[-5::-1,:][0:2000,:][-1::-1,:])

    # fine the start of the cargo container
    startBinCargo = PBAR_Cargo.FindCargoStart(A)
    startBinCargoList.append(startBinCargo)
    datCargoStandardWidth.append(A[(startBinCargo-preCargoTimeSlices):,:][0:numberTimeSlices])
    
del A,bpp,formatt,flag,low1,high1,low2,high2


energyList = []
dat = []
datStandardWidth = []
startBinZspecList = []
for (index, fullFilename) in enumerate(fullFilenameList):
    (a, b) = PBAR_Zspec.ReadZspecBasicScanNumpy(fullFilename)
    energyList.append(a)
    dat.append(b)
#    datStandardWidth.append(b[-1::-1,:,:][0:2000,:,:][-1::-1,:,:])
    startBinZspec = PBAR_Zspec.FindBasicZspecStart(b)
    startBinZspecList.append(startBinZspec)
    datStandardWidth.append(b[(startBinZspec-preCargoTimeSlices):,:][0:numberTimeSlices,:])
    
del a, b


#############################
##  CALCULATE FEATURES

discrim = []
for (index, d) in enumerate(datStandardWidth):
    temp = {}
    energy = energyList[index].astype(float)
    
    temp['count'] = d.sum(2).astype(float)
    
    energyMatrix = np.tile(energy, (d.shape[0], d.shape[1], 1))
    temp['binMean'] = (d * energyMatrix).sum(2) / d.sum(2)
    
    binMeanMatrix = np.tile(temp['binMean'], (d.shape[2],1,1)).transpose((1, 2, 0))
    temp['binSTD'] = np.sqrt( (( (energyMatrix - binMeanMatrix)**2) * d.astype(float) ).sum(2) / d.sum(2).astype(float)  )

    temp['binSTD_binMean'] = temp['binSTD'] / temp['binMean']
    
    # find the bin that is just above 10, 20
    highlowIndex10 = np.argmin( abs(energy - 10) )
    highlowIndex20 = np.argmin( abs(energy - 20) )
    
    temp['multibin_0_10'] = d[:,:,0:highlowIndex10].sum(axis=2)  # bins 0 to 9, so first ten bins
    temp['multibin_10_256'] = d[:,:,highlowIndex10:].sum(axis=2)
    temp['multibin_0_20'] = d[:,:,0:highlowIndex20].sum(axis=2)
    temp['multibin_20_256'] = d[:,:,highlowIndex20:].sum(axis=2)

    temp['multibin_10_ratio'] = temp['multibin_0_10'] / temp['multibin_10_256']
    temp['multibin_20_ratio'] = temp['multibin_0_20'] / temp['multibin_20_256']

    # QUANTILES
    datTempCumSum = d.cumsum(axis = 2).astype(float) # cumulative sum
    datTempCumSum = datTempCumSum / np.tile(datTempCumSum[:,:,-1], (32, 1, 1)).transpose((1, 2, 0)) # normalize the cumulative sum
    qIntervals = [10., 25., 50., 75., 90.]
    for j in range(len(qIntervals)):
        statName = 'q_%02d' % qIntervals[j]# apparently needs a temporary variable
        temp[statName] = energy[np.argmin(abs(datTempCumSum - (qIntervals[j]/100)), axis = 2)]
    temp['q_range'] = temp['q_75']- temp['q_25']
    temp['q_range_ratio'] = temp['q_range'] / temp['q_50']
    temp['binMean_q50'] = temp['binMean'] - temp['q_50']
    discrim.append(temp)

# make truncated image and mask
zspecTMask = []

# http://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp2d.html
#>>> from scipy import interpolate
#>>> x = np.arange(-5.01, 5.01, 0.25)
#>>> y = np.arange(-5.01, 5.01, 0.25)
#>>> xx, yy = np.meshgrid(x, y)
#>>> z = np.sin(xx**2+yy**2)
#>>> f = interpolate.interp2d(x, y, z, kind='cubic')

datCargoZspecStandardWidth = []

for i in xrange(len(dat)):
    mask = np.zeros(datStandardWidth[i].shape)
    numpixels_y_rad = datCargoStandardWidth[i].shape[1]
    numpixels_y_zspec = datStandardWidth[i].shape[1]
    
    y_zspec = np.arange(0,numpixels_y_zspec)
    # the coordinates in radiography units
    y_rad0 = np.arange(0,numpixels_y_rad)
    # the coordinates of the zspec detectors in radiography units
    y_rad = y_zspec * float(numpixels_y_rad) / float(numpixels_y_zspec) - 4.0
    
    x = np.arange(0,datCargoStandardWidth[i].shape[0])
    #xx, yy = np.meshgrid(x, y_rad)
    # the dimensions of  datCargoStandardWidth[i] is ~2000 x 548
    f = interpolate.interp2d(y_rad0, x, datCargoStandardWidth[i], kind='cubic')
    temp = f(y_rad, x)
    # remove negative values.
    temp[temp<0] = 0.0
    datCargoZspecStandardWidth.append(f(y_rad, x))
del f

# Calculate the dist, linear combination
wBest = np.array([-2.2806503 ,  2.26895887,  6.7898522 ])
for index in xrange(len(discrim)):
    # make features array (num features X time slices X number of detectors)
    featuresMatrix = np.zeros((3, discrim[index]['binMean'].shape[0], discrim[index]['binMean'].shape[1]))
    featuresMatrix[0,:,:] = discrim[index]['binMean']
    featuresMatrix[1,:,:] = discrim[index]['binSTD']
    featuresMatrix[2,:,:] = discrim[index]['binMean_q50']
    
    dist0 = featuresMatrix[0,:,:] *wBest[0] + featuresMatrix[1,:,:] *wBest[1] + featuresMatrix[2,:,:] *wBest[2]
    # clean up bad values
    cutt = np.isnan(dist0)
    dist0[cutt] = 0.0
    
    discrim[index]['dist0'] = dist0


# Cut in radiography image of standard size
threshold = 1.5e7 # only consider parts of the zspec image that have rad below this value
cargoZspecStandardWidthMask = []
for i in xrange(len(datCargoZspecStandardWidth)):
    cargoZspecStandardWidthMask.append(datCargoZspecStandardWidth[i] <= threshold)


# Clean up radiograph mask
cargoZspecStandardWidthCleanedMask = []
for i in xrange(len(datCargoZspecStandardWidth)):
    temp = cargoZspecStandardWidthMask[i]
    # do some stuff
    temp = ndimage.binary_opening(temp)
    cargoZspecStandardWidthCleanedMask.append(temp)


# DISPLAY radiograph scaled to zspec standard width
plt.figure()
plt.grid()
index = 0
plt.imshow(np.log(datCargoZspecStandardWidth[index].T+1.0), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargoZspecStandardWidth[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2200))


# DISPLAY radiograph scaled to zspec standard width
plt.figure()
plt.grid()
index = 1
plt.imshow(datCargoZspecStandardWidth[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r, vmax = 8e8)
#plt.imshow(datCargoZspecStandardWidth[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2200))


# DISPLAY radiograph mask
plt.figure()
plt.grid()
index = 0
plt.imshow(cargoZspecStandardWidthMask[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2200))


# DISPLAY Cleaned radiograph mask
plt.figure()
plt.grid()
index = 0
plt.imshow(cargoZspecStandardWidthCleanedMask[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2200))



# IMAGE - linear combination 
plt.figure()
plt.grid()
index = 0
#t = plt.imshow(dist0[:,goodZspecIndices].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
t = plt.imshow(dist0.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)

#t = plt.imshow(dist0[goodZspecIndices,:], interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r, vmin = 100, vmax = 400)
plt.colorbar()
plt.title('Linear Combination')
plt.xlim((0, 2200))


# IMAGE - linear combination with mask
plt.figure()
plt.grid()
index = 0
im = copy.copy(discrim[index]['dist0'].T)
minValue = im.min()
im[~cargoZspecStandardWidthCleanedMask[index].T] = minValue

t = plt.imshow(im, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#t = plt.imshow(dist0[goodZspecIndices,:], interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r, vmin = 100, vmax = 400)
plt.colorbar()
plt.title('Linear Combination')
plt.xlim((0, 2200))


## SCATTER - dist0 vs count
index = 0
# make arrays for counts and dist0, whole image
x1 = discrim[index]['count'][:,goodZspecIndices].flatten()
y1 = discrim[index]['dist0'][:,goodZspecIndices].flatten()

# make arrays for counts and dist0, in mask
im = copy.copy(discrim[index]['count'])
minValue = im.min()
im[~cargoZspecStandardWidthCleanedMask[index]] = minValue
x2 = im[:,goodZspecIndices].flatten()

im = copy.copy(discrim[index]['dist0'])
minValue = im.min()
im[~cargoZspecStandardWidthCleanedMask[index]] = minValue
y2 = im[:,goodZspecIndices].flatten()

obj = []

objtemp = {}
objtemp['xbounds'] = np.array((1118, 1176))
objtemp['ybounds'] = np.array((65, 75))
objtemp['name'] = 'S5'
objtemp['cargo'] = 'sand'
obj.append(objtemp)

objtemp = {}
objtemp['xbounds'] = np.array((1618, 1659))
objtemp['ybounds'] = np.array((82, 85))
objtemp['name'] = 'S5'
objtemp['cargo'] = 'sand'
obj.append(objtemp)


plt.figure()
plt.grid()
plt.plot(x1, y1, '.k', label = 'All Bins', alpha = 0.15, markersize = 10)
plt.plot(x2, y2, '.b', label = 'Mask Bins', alpha = 0.15, markersize = 10)



# HISTOGRAM OF THE dist with and without the mask
index = 0
plt.figure()
plt.grid()
index = 0

binedges = np.arange(-180, 180, 5)
bincenters = (binedges[0:-1] + binedges[1:])/2.0
h1, temp = np.histogram(discrim[index]['dist0'][:,goodZspecIndices].flatten(), binedges)
#h1, temp = np.histogram(discrim[index]['dist0'].flatten(), binedges)

im = copy.copy(discrim[index]['dist0'])
minValue = im.min()
im[~cargoZspecStandardWidthCleanedMask[index]] = minValue
#h2, temp = np.histogram(im.flatten(), binedges)
h2, temp = np.histogram(im[:,goodZspecIndices].flatten(), binedges)

plot(bincenters, h1, label = 'All Bins')
plot(bincenters, h2, label = 'Mask Bins')

plt.title('Linear Combination')
plt.yscale('log')
plt.legend()




#######################################
# MISCELLANEOUS DIAGNOSTIC PLOTS
########################################

# PLOT IMAGE OF ~ STANDARD SIZE BY REVERSING AND OFFSETTING.
plt.figure()
plt.grid()

index = 1
#plt.imshow(datCargo[index].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
# diagnostic plots


# Plot counts vs detector number to see which detectors are bad
plt.figure()
plt.grid()

plt.plot(np.arange(1,137), dat[0][0:100,:,:].sum(2).sum(0))

plt.xlabel('Detector Number')
plt.ylabel('Counts')


# Comparing cargo image and zspec image
# plot of the counts vs time slice
index = 0

plt.figure()
plt.grid()

#plt.plot(dat[index][:,goodZspecIndices,:].sum(2).sum(1) / dat[index][:,goodZspecIndices,:].sum(2).sum(1).mean(), label = 'Zspec')
plt.plot(dat[index][:,25:28,:].astype(float).sum(2).sum(1) / dat[index][:,25:28,:].astype(float).sum(2).sum(1).mean(), label = 'Zspec')
plt.plot(datCargo[index][:,90:110].astype(float).sum(1) / datCargo[index][:,90:110].astype(float).sum(1).mean(), label = 'Radiograph')

plt.xlabel('Slice')
plt.ylabel('Signal')
plt.legend()




# PLOT unmodified cargo image
index = 1

plt.figure()
plt.grid()

#plt.plot(dat[index][:,goodZspecIndices,:].sum(2).sum(1) / dat[index][:,goodZspecIndices,:].sum(2).sum(1).mean(), label = 'Zspec')
plt.imshow(datCargo[index].astype(float).T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r, vmax = 5e8)
marker = markerList[index]
for m in marker:
    xCenter = m['x']
    yCenter = m['y']
    x = np.array((m['left'], m['right']))
    y = np.array((m['top'], m['bottom']))
    plt.plot(xCenter, yCenter, 'xr', markersize = 15)
    plt.plot((x[0], x[1], x[1], x[0], x[0]), (y[0], y[0], y[1], y[1], y[0]), '-r')    
plt.xlabel('Slice')
plt.ylabel('Signal')
plt.legend()




# PLOT BOTH IMAGES WITH BEST OVERLAP POSSIBLE

threshold = 1.5e9

plt.figure()
plt.grid()

index = 1
plt.imshow(log(datCargoStandardWidth[index].T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#plt.imshow(datCargo[index][-5::-1,:].T < threshold, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(filenameCargoList[index])
plt.xlim((0, 2200))


# zspec image
plt.figure()
plt.grid()

discrimName = 'binSTD_binMean'
discrimName = 'binSTD'
discrimName = 'count'
#discrimName = 'multibin_20_ratio'
index = 1
plt.imshow(discrim[index][discrimName].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r, vmin = 0, vmax = 500)
#plt.imshow(discrim[index][discrimName][-1::-1,1:].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
plt.colorbar()
plt.title(discrimName)
plt.xlim((0, 2200))



# figure out threshold to use to remove the high transmission parts of the image

# Make an animated GIF
from PIL import Image, ImageSequence
from images2gif import writeGif

import sys, os
frames = []
frames.append(Image.open(r'C:\Users\jkwong\Documents\Work\PBAR\Plots\20130919\4939.png'))
frames.append(Image.open(r'C:\Users\jkwong\Documents\Work\PBAR\Plots\20130919\PBAR-20130902007.png'))

writeGif(r'C:\Users\jkwong\Documents\Work\PBAR\Plots\20130919\4939_PBAR-20130902007.gif', frames, duration=2.0, loops = 0, dither=1)


